import cv2
import numpy as np
from utils.utils import decode, non_max_suppression, draw_image, read_label_names

# 加载模型
model = cv2.dnn.readNetFromONNX('models.onnx')

# 定义类别标签
class_names = read_label_names('./config/class_names.txt')
class_nums = len(class_names)
# 定义置信度阈值和NMS阈值
conf_threshold = 0.5
nms_threshold = 0.5

# 加载图像
img = cv2.imread('img/street.jpg')
img = cv2.resize(img, dsize=None, fx=0.5, fy=0.5)

image_shape = np.array(np.shape(img)[0:2])

# 将图像转换为blob格式
blob = cv2.dnn.blobFromImage(img, 1 / 255, (416, 416), swapRB=True, crop=False)

# 将blob输入到模型中获取输出
model.setInput(blob)

outputs = model.forward(model.getUnconnectedOutLayersNames())

result = decode(outputs, num_classes=class_nums)
results = non_max_suppression(result, num_classes=class_nums, image_shape=image_shape)

img = draw_image(img, results, class_names)
cv2.imshow(' success', img)
cv2.waitKey(0)
